from sklearn.datasets import load_boston
import torch
import torch.nn as nn
import torchvision


class LineModel(nn.Module):
    def __init__(self, ndim):
        super(LineModel, self).__init__()
        self.ndim = ndim
        self.weight = nn.Parameter(torch.randon(ndim, 1))
        self.bias = nn.Parameter(torch.randn(1))

    def forward(self, x):
        return x.mm(self.weight) + self.bias


if __name__ == '__main__':
    boston = load_boston()
    torchvision.models.resnet18
    lm = LineModel(13)
    criterion = nn.MSELoss()
    optim = torch.optim.SGD(lm.parameters(), lr=1e-6)
    data = torch.tensor(boston["data"], requires_grad=True)
    target = torch.tensor(boston["target"], dtype=torch.float32)
    step = 0

    for step in range(1000):
        predict = lm(data)
        loss = criterion(predict, target)
        if step and step % 100:
            print("Loss:{::.3f}".format(loss.item()))
        optim.zero_grad()
        loss.backward()
        optim.step()
